// Copyright 2022 Chen Jun
// Licensed under the MIT License.

#include <cv_bridge/cv_bridge.h>
#include <rmw/qos_profiles.h>
#include <tf2/LinearMath/Matrix3x3.h>
#include <tf2/convert.h>

#include <ament_index_cpp/get_package_share_directory.hpp>
#include <image_transport/image_transport.hpp>
#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <rclcpp/duration.hpp>
#include <rclcpp/qos.hpp>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>

// STD
#include <algorithm>
#include <map>
#include <memory>
#include <string>
#include <vector>
//AI algorithm

#include "armor_detector/armor.hpp"
#include "armor_detector/detector_node.hpp"

namespace rm_auto_aim
{
ArmorDetectorNode::ArmorDetectorNode(const rclcpp::NodeOptions & options) // 定义构造函数，传入ROS节点选项
: Node("armor_detector", options)// 调用Node类构造函数，设置节点名称和选项
{
  RCLCPP_INFO(this->get_logger(), "Starting DetectorNode!");// 输出节点启动信息

  // Detector
  detector_ = initDetector();// 初始化装甲板检测器

  //1. Armors Publisher--创建一个发布者，用于发布检测到的装甲板信息
  armors_pub_ = this->create_publisher<auto_aim_interfaces::msg::Armors>(
    "/detector/armors", rclcpp::SensorDataQoS());

  //2.Visualization Marker Publisher--创建一个发布者，用于发布可视化标记信息
  // See http://wiki.ros.org/rviz/DisplayTypes/Marker
  armor_marker_.ns = "armors";
  armor_marker_.action = visualization_msgs::msg::Marker::ADD;
  armor_marker_.type = visualization_msgs::msg::Marker::CUBE;
  armor_marker_.scale.x = 0.05;
  armor_marker_.scale.z = 0.125;
  armor_marker_.color.a = 1.0;
  armor_marker_.color.g = 0.5;
  armor_marker_.color.b = 1.0;
  armor_marker_.lifetime = rclcpp::Duration::from_seconds(0.1);

  text_marker_.ns = "classification";
  text_marker_.action = visualization_msgs::msg::Marker::ADD;
  text_marker_.type = visualization_msgs::msg::Marker::TEXT_VIEW_FACING;
  text_marker_.scale.z = 0.1;
  text_marker_.color.a = 1.0;
  text_marker_.color.r = 1.0;
  text_marker_.color.g = 1.0;
  text_marker_.color.b = 1.0;
  text_marker_.lifetime = rclcpp::Duration::from_seconds(0.1);

  marker_pub_ =
    this->create_publisher<visualization_msgs::msg::MarkerArray>("/detector/marker", 10);

  // 3.Debug Publishers--创建调试信息的发布者
  debug_ = this->declare_parameter("debug", false);
  if (debug_) {
    createDebugPublishers();
  }
  //4. Debug param change moniter--创建一个ROS参数事件处理器：监视debug的参数
  debug_param_sub_ = std::make_shared<rclcpp::ParameterEventHandler>(this);
  debug_cb_handle_ =
    debug_param_sub_->add_parameter_callback("debug", [this](const rclcpp::Parameter & p) {
      debug_ = p.as_bool();
      debug_ ? createDebugPublishers() : destroyDebugPublishers();// 根据调试参数的值，创建或销毁调试信息的发布者
    });
  /*5.创建了一个名为cam_info_sub_的订阅者对象，用于订阅名为/camera_info的ROS话题.
@func:在接收到消息时调用回调函数。该回调函数会读取相机参数信息，并保存到节点的成员变量中，以便后续使用*/
  cam_info_sub_ = this->create_subscription<sensor_msgs::msg::CameraInfo>(
    "/camera_info", rclcpp::SensorDataQoS(),
    [this](sensor_msgs::msg::CameraInfo::ConstSharedPtr camera_info) {
      cam_center_ = cv::Point2f(camera_info->k[2], camera_info->k[5]);// 读取相机参数，用于计算装甲板的相对位置和角度
      cam_info_ = std::make_shared<sensor_msgs::msg::CameraInfo>(*camera_info);// 保存相机参数
      pnp_solver_ = std::make_unique<PnPSolver>(camera_info->k, camera_info->d); // 创建一个PnP求解器，用于计算装甲板的相对位置和角度
      cam_info_sub_.reset();
    });
    /*6.创建一个订阅者，用于接收图像信息，并绑定回调函数imageCallback
    @订阅名为"/image_raw"
    @消息类型为<sensor_msgs::msg::Image>
    @使用了rclcpp::SensorDataQoS()作为其质量保证，表示它只接收最新的传感器数据
    @当有新的图像消息发布到"/image_raw"话题时，会调用ArmorDetectorNode类中的imageCallback函数来处理接收到的图像数据
    */ 
  img_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
    "/image_raw", rclcpp::SensorDataQoS(),
    std::bind(&ArmorDetectorNode::imageCallback, this, std::placeholders::_1));
    //调用回调函数，而回调函数中包含了装甲板检测代码（关键部分1）--如何传入img给imagecallback函数呢？（问题1）
}
/*imageCallback:功能是处理接收到的图像消息，进行装甲板检测并发布检测结果和可视化标记信息。*/
void ArmorDetectorNode::imageCallback(const sensor_msgs::msg::Image::ConstSharedPtr img_msg)
{
  //对接收到的图像进行装甲板检测，得到检测到的装甲板信息（关键部分2）
  auto armors = detectArmors(img_msg);
  //根据装甲板的位置和姿态信息，构建装甲板的消息并填充到armors_msg_中
  if (pnp_solver_ != nullptr) {
    armors_msg_.header = armor_marker_.header = text_marker_.header = img_msg->header;
    armors_msg_.armors.clear();
    marker_array_.markers.clear();
    armor_marker_.id = 0;
    text_marker_.id = 0;

    auto_aim_interfaces::msg::Armor armor_msg;
    for (const auto & armor : armors) {
      cv::Mat rvec, tvec;
      bool success = pnp_solver_->solvePnP(armor, rvec, tvec);
      if (success) {
        // Fill basic info
        armor_msg.type = ARMOR_TYPE_STR[static_cast<int>(armor.type)];
        armor_msg.number = armor.number;

        // Fill pose
        armor_msg.pose.position.x = tvec.at<double>(0);
        armor_msg.pose.position.y = tvec.at<double>(1);
        armor_msg.pose.position.z = tvec.at<double>(2);
        // rvec to 3x3 rotation matrix
        cv::Mat rotation_matrix;
        cv::Rodrigues(rvec, rotation_matrix);
        // rotation matrix to quaternion
        tf2::Matrix3x3 tf2_rotation_matrix(
          rotation_matrix.at<double>(0, 0), rotation_matrix.at<double>(0, 1),
          rotation_matrix.at<double>(0, 2), rotation_matrix.at<double>(1, 0),
          rotation_matrix.at<double>(1, 1), rotation_matrix.at<double>(1, 2),
          rotation_matrix.at<double>(2, 0), rotation_matrix.at<double>(2, 1),
          rotation_matrix.at<double>(2, 2));
        tf2::Quaternion tf2_q;
        tf2_rotation_matrix.getRotation(tf2_q);
        armor_msg.pose.orientation = tf2::toMsg(tf2_q);

        // Fill the distance to image center
        armor_msg.distance_to_image_center = pnp_solver_->calculateDistanceToCenter(armor.center);
        //将装甲板的可视化标记信息填充到marker_array_中
        // Fill the markers
        armor_marker_.id++;
        armor_marker_.scale.y = armor.type == ArmorType::SMALL ? 0.135 : 0.23;
        armor_marker_.pose = armor_msg.pose;
        text_marker_.id++;
        text_marker_.pose.position = armor_msg.pose.position;
        text_marker_.pose.position.y -= 0.1;
        text_marker_.text = armor.classfication_result;
        armors_msg_.armors.emplace_back(armor_msg);
        marker_array_.markers.emplace_back(armor_marker_);
        marker_array_.markers.emplace_back(text_marker_);
      } else {
        RCLCPP_WARN(this->get_logger(), "PnP failed!");
      }
    }

    // Publishing detected armors--通过armors_pub_发布检测到的装甲板信息
    armors_pub_->publish(armors_msg_);

    // Publishing marker--通过publishMarkers函数发布装甲板的可视化标记信息。
    publishMarkers();
  }
}
/*读取ROS2参数服务器中的参数，包括二值化阈值、检测颜色、灯条参数、装甲板参数、分类器模型路径、分类器阈值等参数。*/
std::unique_ptr<Detector> ArmorDetectorNode::initDetector()
{
  rcl_interfaces::msg::ParameterDescriptor param_desc;//声明一个参数描述符对象，用于描述后续的变量的类型和范围
  param_desc.integer_range.resize(1);//设置参数描述符中整数范围的大小为1，表示只有一个范围。
  param_desc.integer_range[0].step = 1;//步长为1
  param_desc.integer_range[0].from_value = 0;//整数范围起始值为0
  param_desc.integer_range[0].to_value = 255;//整数范围结束值为255
  int binary_thres = declare_parameter("binary_thres", 160, param_desc);//从ROS2参数服务器中获取名为"binary_thres"的参数，如果没有找到该参数，则使用默认值160。param_desc参数用于指定该参数的描述符，包括类型和范围。
  /* 为参数'detect_color'定义描述*/
  param_desc.description = "0-RED, 1-BLUE";
  param_desc.integer_range[0].from_value = 0;// 将整数范围的最小值设置为0
  param_desc.integer_range[0].to_value = 1;// 将整数范围的最大值设置为1
  auto detect_color = declare_parameter("detect_color", RED, param_desc);// 使用提供的描述和范围声明名为'detect_color'的参数，并设定默认值为RED
  //描述灯条灯条的长度与宽度的比值，灯条是指在图像中表示一个物体（如车辆或机器人）的矩形区域；这个比率可以用来过滤掉那些宽度相对于长度太小的灯条，从而只保留那些比例合适的灯条进行后续处理
  Detector::LightParams l_params = {
    .min_ratio = declare_parameter("light.min_ratio", 0.1),
    .max_ratio = declare_parameter("light.max_ratio", 0.4),
    // 将最大角度设定为从参数'light.max_angle'获得的值
    .max_angle = declare_parameter("light.max_angle", 40.0)};
  //装甲板参数
  Detector::ArmorParams a_params = {
    // 将最小灯条亮度比率设定为从参数'armor.min_light_ratio'获得的值
    /*，装甲板通常由两个灯条组成，这两个灯条之间的距离、角度和亮度（或颜色）等方面的特征可以用于判断是否是一个有效的装甲板。
    而灯条的亮度（或颜色）与周围环境的亮度（或颜色）之间的比率可以用来提高装甲板的识别率。
    因此，通过设置一个最小的灯条亮度比率阈值，来过滤掉那些与周围环境亮度（或颜色）相似的区域，从而只保留那些比例合适的灯条，用于后续处理*/
    .min_light_ratio = declare_parameter("armor.min_light_ratio", 0.7),
    // 设置小目标中心之间的最小距离，用于过滤掉太小的目标。
    .min_small_center_distance = declare_parameter("armor.min_small_center_distance", 0.8),
    //描述小目标之间的最大中心间距，用于过滤掉太远或太近的目标。
    .max_small_center_distance = declare_parameter("armor.max_small_center_distance", 3.2),
    .min_large_center_distance = declare_parameter("armor.min_large_center_distance", 3.2),
    .max_large_center_distance = declare_parameter("armor.max_large_center_distance", 5.5),
    //两个装甲板之间的最大角度
    .max_angle = declare_parameter("armor.max_angle", 35.0)};
  //创建一个装甲板检测器对象，并将读取的参数传递给检测器对象。
  auto detector = std::make_unique<Detector>(binary_thres, detect_color, l_params, a_params);

  // Init classifier创建一个数字分类器对象，并将读取的分类器模型路径、标签路径、阈值等参数传递给分类器对象。
  auto pkg_path = ament_index_cpp::get_package_share_directory("armor_detector");
  auto model_path = pkg_path + "/model/mlp.onnx";
  auto label_path = pkg_path + "/model/label.txt";
  double threshold = this->declare_parameter("classifier_threshold", 0.7);
  std::vector<std::string> ignore_classes =
    this->declare_parameter("ignore_classes", std::vector<std::string>{"negative"});
  detector->classifier =
    std::make_unique<NumberClassifier>(model_path, label_path, threshold, ignore_classes);

  return detector;
}

std::vector<Armor> ArmorDetectorNode::detectArmors(
  const sensor_msgs::msg::Image::ConstSharedPtr & img_msg)
{
  // 将ROS图像转换为OpenCV的cv::Mat格式
  auto img = cv_bridge::toCvShare(img_msg, "rgb8")->image;
  // 更新参数
  detector_->binary_thres = get_parameter("binary_thres").as_int();
  detector_->detect_color = get_parameter("detect_color").as_int();
  detector_->classifier->threshold = get_parameter("classifier_threshold").as_double();
  // 使用检测器对图像进行装甲板检测
  auto armors = detector_->detect(img);
  // 计算识别耗时
  auto final_time = this->now();
  auto latency = (final_time - img_msg->header.stamp).seconds() * 1000;
  RCLCPP_DEBUG_STREAM(this->get_logger(), "Latency: " << latency << "ms");
  // Publish debug info发布调试信息
  if (debug_) {
    // 发布二值化图像
    binary_img_pub_.publish(
      cv_bridge::CvImage(img_msg->header, "mono8", detector_->binary_img).toImageMsg());
    // 根据x坐标对灯条和装甲板数据进行排序--按照它们在图像中的 x 坐标从小到大排列。
    std::sort(
      detector_->debug_lights.data.begin(), detector_->debug_lights.data.end(),
      [](const auto & l1, const auto & l2) { return l1.center_x < l2.center_x; });
    std::sort(
      detector_->debug_armors.data.begin(), detector_->debug_armors.data.end(),
      [](const auto & a1, const auto & a2) { return a1.center_x < a2.center_x; });
    // 发布灯条和装甲板数据
    lights_data_pub_->publish(detector_->debug_lights);
    armors_data_pub_->publish(detector_->debug_armors);
    // 如果存在装甲板，发布数字图像
    if (!armors.empty()) {
      auto all_num_img = detector_->getAllNumbersImage();
      number_img_pub_.publish(
        *cv_bridge::CvImage(img_msg->header, "mono8", all_num_img).toImageMsg());
    }
    // 绘制检测结果
    detector_->drawResults(img);
    //  在图像上绘制相机中心点
    cv::circle(img, cam_center_, 5, cv::Scalar(255, 0, 0), 2);
    // 在图像上绘制识别耗时
    std::stringstream latency_ss;
    latency_ss << "Latency: " << std::fixed << std::setprecision(2) << latency << "ms";
    auto latency_s = latency_ss.str();
    cv::putText(
      img, latency_s, cv::Point(10, 30), cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(0, 255, 0), 2);
    // 发布最终结果图像
    result_img_pub_.publish(cv_bridge::CvImage(img_msg->header, "rgb8", img).toImageMsg());
  }
  // 返回识别到的装甲板
  return armors;
}

//创建用于调试目的的发布者（publishers）
void ArmorDetectorNode::createDebugPublishers()
{
  //灯条的调试信息
  lights_data_pub_ =
    this->create_publisher<auto_aim_interfaces::msg::DebugLights>("/detector/debug_lights", 10);
  //甲板的调试信息
  armors_data_pub_ =
    this->create_publisher<auto_aim_interfaces::msg::DebugArmors>("/detector/debug_armors", 10);
  //发布二值化图像到话题 /detector/binary_img
  binary_img_pub_ = image_transport::create_publisher(this, "/detector/binary_img");
  //发布数字图像
  number_img_pub_ = image_transport::create_publisher(this, "/detector/number_img");
  //发布结果图像
  result_img_pub_ = image_transport::create_publisher(this, "/detector/result_img");
}
//作用是销毁调试发布者（publishers）
void ArmorDetectorNode::destroyDebugPublishers()
{
  //将 lights_data_pub_ 和 armors_data_pub_ 的指针设置为空指针，即释放相关资源。
  lights_data_pub_.reset();
  armors_data_pub_.reset();
  //调用 binary_img_pub_、number_img_pub_ 和 result_img_pub_ 的 shutdown() 函数，关闭对应的发布者。
  binary_img_pub_.shutdown();
  number_img_pub_.shutdown();
  result_img_pub_.shutdown();
}
//作用是发布标记信息，用于可视化装甲板的位置--通过调用这个函数，可以将装甲板的位置标记进行可视化，方便在可视化工具中查看装甲板的位置信息。
void ArmorDetectorNode::publishMarkers()
{
  using Marker = visualization_msgs::msg::Marker;
  //判断是否为空--为空则代表没有检测到装甲板，则执行delete行为；如果不为空，则执行add行为。
  armor_marker_.action = armors_msg_.armors.empty() ? Marker::DELETE : Marker::ADD;
  marker_array_.markers.emplace_back(armor_marker_);//将 armor_marker_ 添加到 marker_array_ 中；marker_array用于存储多个标记信息
  marker_pub_->publish(marker_array_);//通过 marker_pub_ 发布 marker_array_，以将标记信息发布到相应的话题。
}

}  // namespace rm_auto_aim

#include "rclcpp_components/register_node_macro.hpp"

// Register the component with class_loader.
// This acts as a sort of entry point, allowing the component to be discoverable when its library
// is being loaded into a running process.
RCLCPP_COMPONENTS_REGISTER_NODE(rm_auto_aim::ArmorDetectorNode)
